import io
from PIL import Image
from docx import Document


# NOTE:
# DocxExtractor 可用, 但是提取图片不能分到段落, 先放着吧
# extract_targets_docx 及 extract_and_show_images 相当于
# DocxExtractor.extract_paragraphs 及 DocxExtractor.extract_paragraphs 的函数调用


class DocxExtractor:
    def __init__(self, path_docx):
        self.path_docx = path_docx
        self.matched_paragraphs = {     # 保存段落
            "3集成电路设计实验结果": {
                "word": [],
            },
            "4集成电路制造实验结果": {
                "word": [],
            },
            "5集成电路封装与测试实验结果": {
                "word": [],
            },
            "6实验结果分析": {
                "word": [],
            },
            "7实验结论": {
                "word": [],
            }
        }
        self.match_images = []      # 保存图片
        self.doc = None
        self.current_title = None

    def load_document(self):
        """加载文档"""
        self.doc = Document(self.path_docx)

    def extract_paragraphs(self):
        """提取匹配的段落"""
        if self.doc is None:
            self.load_document()
        # 提取段落
        for para in self.doc.paragraphs:
            # 检查段落的样式是否为“标题1”
            if para.style.name == 'Heading 1' and para.text.replace(".", "").replace(" ", "") in self.matched_paragraphs:
                self.current_title = para.text
            elif self.current_title is not None:
                # 如果当前段落不是“标题1”，并且当前有正在处理的标题
                if para.style.name != 'Heading 1':
                    # 将段落文字添加到相应标题的word列表中
                    self.matched_paragraphs[self.current_title]["word"].append(para.text)
                else:
                    # 如果遇到了新的“标题1”，重置current_title
                    self.current_title = para.text if para.text in self.matched_paragraphs else None

    # FIXME: 提取图片到段落
    def extract_images(self):
        """从文档中提取图片数据"""
        if self.doc is None:
            self.load_document()
        #
        image_rels = {rel.rId: rel.target_part.blob for rel in self.doc.part.rels.values() if "image" in rel.target_ref}
        #
        for inline_shape in self.doc.inline_shapes:
            if inline_shape.type.value == 3:  # 3代表图片
                rId = inline_shape._inline.graphic.graphicData.pic.blipFill.blip.embed
                if rId in image_rels:
                    image_data = image_rels[rId]
                    self.match_images.append(image_data)

    def get_matched_paragraphs(self):
        """返回匹配的段落字典"""
        return self.matched_paragraphs

    def get_matched_images(self):
        """返回图片列表"""
        return self.match_images




# ====================================== 供外部调用的方法 =======================================
def extract_targets_docx(path_docx):
    # 用于存储匹配的段落
    dict_matched_paragraphs = {
        "3集成电路设计实验结果": {
            "word": [],
            "image": [],
        },
        "4集成电路制造实验结果": {
            "word": [],
            "image": [],
        },
        "5集成电路封装与测试实验结果": {
            "word": [],
            "image": [],
        },
        "6实验结果分析": {
            "word": [],
            "image": [],
        },
        "7实验结论": {
            "word": [],
            "image": [],
        }
    }
    # 打开文档
    doc = Document(path_docx)
    # 轮询doc.paragraphs
    current_title = None
    for para in doc.paragraphs:
        # 检查段落的样式是否为“标题1”
        if para.style.name == 'Heading 1' and para.text in dict_matched_paragraphs:
            current_title = para.text
        elif current_title is not None:
            # 如果当前段落不是“标题1”，并且当前有正在处理的标题
            if para.style.name != 'Heading 1':
                # 将段落文字添加到相应标题的word列表中
                dict_matched_paragraphs[current_title]["word"].append(para.text)
            else:
                # 如果遇到了新的“标题1”，重置current_title
                current_title = para.text if para.text in dict_matched_paragraphs else None
    return dict_matched_paragraphs


def extract_and_show_images(docx_path):
    match_images = []
    # 打开文档
    doc = Document(docx_path)
    # 用于存储图片关系
    image_rels = {}
    for rel in doc.part.rels.values():
        if "image" in rel.target_ref:
            image_rels[rel.rId] = rel.target_part.blob
    # 遍历所有的内联形状
    for i, inline_shape in enumerate(doc.inline_shapes):
        print(inline_shape.type, inline_shape.type.value)
        if inline_shape.type.value == 3:  # 3代表图片
            # 获取图片所在的段落索引
            # 获取图片的 rId
            rId = inline_shape._inline.graphic.graphicData.pic.blipFill.blip.embed
            # 从关系中获取图片数据
            if rId in image_rels:
                image_data = image_rels[rId]
                # # 使用Pillow打开图片
                # image = Image.open(io.BytesIO(image_data))
                # # 显示图片
                # image.show()
                match_images.append(image_data)
    return match_images




# ====================================== 测试函数 =======================================
# 测试 extract_targets_docx()
def test_1():
    doc_path = f'C://work/【实验室相关】/【2024-10-22】AI判分/10.91.128.2-Tomcat80/202108/mj_2153507_鏈辩亸_20210817182547.docx'
    dict_matched_paragraphs = extract_targets_docx(doc_path)
    for key, paras in dict_matched_paragraphs.items():
        print(key)
        for para in paras['word']:
            print(para)
        print("-------------------------------")


# 测试 DocxExtractor
def test_2():
    # 示例调用
    path_docx = f'C://work/【实验室相关】/【2024-10-22】AI判分/10.91.128.2-Tomcat80/202108/mj_2153507_鏈辩亸_20210817182547.docx'
    extractor = DocxExtractor(path_docx)
    extractor.extract_paragraphs()
    result = extractor.get_matched_paragraphs()
    print(result)


# 测试 extract_and_show_images
def test_3():
    path_docx = f'C://work/【实验室相关】/【2024-10-22】AI判分/10.91.128.2-Tomcat80/202108/mj_2153507_鏈辩亸_20210817182547.docx'
    extract_and_show_images(path_docx)




if __name__=="__main__":
    test_3()